#pragma once

#include <iostream>
#include <vector>
#include <thread>
#include <condition_variable>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <string>

typedef int (*HandleFunction)(int s, std::string cmd);

class TCPServer
{
public:
    TCPServer(uint16_t port, HandleFunction handleFunction)
        : port_(port), stop_(false)
    {
        handleFunction_ = handleFunction;
        // 初始化服务器套接字
        if ((server_fd_ = socket(AF_INET, SOCK_STREAM, 0)) == -1)
        {
            throw std::runtime_error("Could not create socket");
        }

        // 设置地址和端口
        address_.sin_family = AF_INET;
        address_.sin_addr.s_addr = INADDR_ANY;
        address_.sin_port = htons(port_);

        int yes = 1;
        if (setsockopt(server_fd_, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)) == -1)
        {
            perror("setsockopt");
            exit(EXIT_FAILURE);
        }

        // 绑定套接字
        if (bind(server_fd_, (struct sockaddr *)&address_, sizeof(address_)) < 0)
        {
            throw std::runtime_error("Bind failed");
        }

        // 监听套接字
        if (listen(server_fd_, 1) < 0) // 只支持连接一个客户端
        {
            throw std::runtime_error("Listen failed");
        }
    }

    ~TCPServer()
    {
        stop();
        close(server_fd_);
    }

    void start()
    {
        stop_ = false;
        accept_thread_ = std::thread(&TCPServer::acceptConnections, this);
    }

    void stop()
    {
        {
            std::unique_lock<std::mutex> lock(mtx_);
            stop_ = true;
            cond_.notify_all();
        }
        if (accept_thread_.joinable())
        {
            accept_thread_.join();
        }
        for (auto &thread : client_threads_)
        {
            if (thread.joinable())
            {
                thread.join();
            }
        }
        client_threads_.clear();
    }

private:
    void acceptConnections()
    {
        int timeout = 500000; // 设置超时为 5000 毫秒（5 秒）
        int set = 1;

        while (!stop_)
        {
            int client_fd = accept(server_fd_, nullptr, nullptr);
            if (client_fd < 0)
            {
                continue;
            }
            setsockopt(client_fd, SOL_SOCKET, SO_SNDTIMEO, (char *)&timeout, sizeof(timeout));
            setsockopt(client_fd, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(timeout));
            std::thread client_thread(&TCPServer::handleClient, this, client_fd);
            {
                std::lock_guard<std::mutex> lock(mtx_);
                client_threads_.push_back(std::move(client_thread));
            }
        }
    }

    void handleClient(int client_fd)
    {
        char buffer[1024];
        uint32_t remain, package_len, len, i;
        int ret = 1;

        while (ret)
        {
            // 读取从客户端发送的数据
            len = read(client_fd, buffer, sizeof(uint32_t));
            if (len == 0)
            {
                std::cout << "client closed " << std::endl;
                return;
            }

            package_len = *(uint32_t *)&buffer[0];
            remain = package_len;
            i = 0;
            do
            {
                len = read(client_fd, buffer + i, package_len - i);
                i += len;
            } while (i < package_len);

            buffer[i + 1] = '\0'; // 结束符
            std::string cmd(buffer);
            ret = handleFunction_(client_fd, cmd);

            // 发送数据到客户端
            // send(client_fd, buffer, package_len, 0);
        }
        close(client_fd);
    }

private:
    HandleFunction handleFunction_;
    uint16_t port_;
    int server_fd_;
    struct sockaddr_in address_;
    std::vector<std::thread> client_threads_;
    std::thread accept_thread_;
    std::mutex mtx_;
    std::condition_variable cond_;
    bool stop_;
};
